from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_ollama import OllamaEmbeddings

# 在bash中还是需要安装faiss的：pip install faiss-gpu
from langchain_community.vectorstores import FAISS

# 使用ollama中的model作为嵌入式模型，提供给vector_store进行使用
embeddings = OllamaEmbeddings(model="deepseek-r1:7b", base_url="http://localhost:11434")

# 测试pdf数据，读取操作
file_path = "./test_pdf/nke-10k-2023.pdf"
loader = PyPDFLoader(file_path)
docs = loader.load()

# 测试验证
# print(f"{docs[10].page_content[:200]}\n")
# print(docs[10].metadata)

# 通过RecursiveCharacterTextSplitter进行text分词测试
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=1000, chunk_overlap=200, add_start_index=True
)
all_splits = text_splitter.split_documents(docs)

# 测试验证
# print("all_splits[0]:"+str(all_splits[0]))
# print("all_splits[1]:"+str(all_splits[1]))

# 使用嵌入式模型验证查询能力
vector_1 = embeddings.embed_query(all_splits[0].page_content)
vector_2 = embeddings.embed_query(all_splits[1].page_content)

assert len(vector_1) == len(vector_2)
print(f"Generated vectors of length {len(vector_1)}\n")
print(vector_1[:10])

# 从文档创建 FAISS 索引，并保存到本地
# vector_store = FAISS.from_documents(all_splits, embeddings)
# vector_store.save_local("./test_vsdata/","index.bin")

# 加载本地存储的 FAISS 向量库（假设你之前已保存到 "faiss_index" 目录）
# 为了安全考虑，默认情况下禁止加载可能存在风险的 pickle 文件。
# 如果你确认文件的来源是可信的，可以通过设置 allow_dangerous_deserialization=True 来允许反序列化。
vector_store = FAISS.load_local("./test_vsdata/", embeddings, 
                                index_name="index.bin",
                                allow_dangerous_deserialization=True)

# 进行相似性搜索，假设输入查询文本
results = vector_store.similarity_search(
    "How many distribution centers does Nike have in the US?"
)
print(results[0])